# app.py
from flask import Flask, request, jsonify
from flask_socketio import SocketIO, emit, join_room
from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity, verify_jwt_in_request
from flask_sqlalchemy import SQLAlchemy
from flask_jwt_extended.exceptions import NoAuthorizationError
import datetime
import socket

app = Flask(__name__)
app.config['SECRET_KEY'] = 'zCH33JZHWyaG4Q7yx7P3'      # 替换成你自己的 SECRET_KEY
app.config['JWT_SECRET_KEY'] = 'zCH33JZHWyaG4Q7yx7P3'  # 替换成你自己的 JWT_SECRET_KEY
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///chat.db'
app.config['JWT_ACCESS_TOKEN_EXPIRES'] = datetime.timedelta(days=1)

db = SQLAlchemy(app)
jwt = JWTManager(app)
socketio = SocketIO(app, cors_allowed_origins="*")


class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password = db.Column(db.String(120), nullable=False)


class Message(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    content = db.Column(db.Text, nullable=False)
    sender = db.Column(db.String(80), nullable=False)
    receiver = db.Column(db.String(80), nullable=False)
    timestamp = db.Column(
        db.DateTime,
        default=lambda: datetime.datetime.now(datetime.timezone(datetime.timedelta(hours=8)))
    )


# 确保在应用上下文中创建表
# with app.app_context():
#     db.create_all()
#     print("Database tables created successfully")


@app.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    if User.query.filter_by(username=data['username']).first():
        return jsonify({"msg": "用户名已存在"}), 400

    new_user = User(username=data['username'], password=data['password'])
    db.session.add(new_user)
    db.session.commit()
    return jsonify({"msg": "注册成功"}), 201


@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    user = User.query.filter_by(username=data['username'], password=data['password']).first()
    if not user:
        return jsonify({"msg": "用户名或密码错误"}), 401

    access_token = create_access_token(identity=user.username)
    return jsonify(access_token=access_token), 200


@app.route('/users', methods=['GET'])
@jwt_required()
def get_users():
    users = User.query.all()
    return jsonify([{'username': user.username} for user in users])


@app.route('/messages', methods=['GET'])
@jwt_required()
def get_messages():
    current_user = get_jwt_identity()
    target_user = request.args.get('with')

    # 获取双方之间的所有消息
    messages = Message.query.filter(
        ((Message.sender == current_user) & (Message.receiver == target_user)) |
        ((Message.sender == target_user) & (Message.receiver == current_user))
    ).order_by(Message.timestamp.asc()).all()

    return jsonify([{
        'id': msg.id,
        'sender': msg.sender,
        'receiver': msg.receiver,
        'content': msg.content,
        'timestamp': msg.timestamp.isoformat()
    } for msg in messages])


# 全局在线用户字典
online_users = {}


@socketio.on('connect')
@jwt_required()
def handle_connect():
    username = get_jwt_identity()
    print(f"用户 {username} 已连接")
    join_room(username)

    # 更新在线状态
    online_users[username] = True

    # 广播用户上线状态
    emit('user_status', {'username': username, 'online': True}, broadcast=True)

    # 发送当前所有在线用户状态给新连接的用户
    emit('initial_online_status', {
        'online_users': [user for user, online in online_users.items() if online]
    }, room=username)


@socketio.on('disconnect')
@jwt_required(optional=True)
def handle_disconnect():
    username = get_jwt_identity()
    print(f"用户 {username} 已断开")

    # 更新在线状态
    online_users[username] = False

    # 广播用户下线状态
    emit('user_status', {'username': username, 'online': False}, broadcast=True)


# 添加获取在线用户列表的路由
@app.route('/online_users', methods=['GET'])
@jwt_required()
def get_online_users():
    return jsonify({
        'online_users': [user for user, online in online_users.items() if online]
    }), 200


# 添加心跳检测路由
@app.route('/heartbeat', methods=['POST'])
@jwt_required()
def heartbeat():
    username = get_jwt_identity()
    online_users[username] = True
    return jsonify({'status': 'ok'}), 200


# 添加加入房间的事件处理
@socketio.on('join_room')
@jwt_required()
def handle_join_room(room):
    username = get_jwt_identity()
    print(f"用户 {username} 加入房间 {room}")
    join_room(room)


@socketio.on('send_message')
@jwt_required()
def handle_message(data):
    sender = get_jwt_identity()
    receiver = data['receiver']

    # 保存消息到数据库
    message = Message(
        content=data['content'],
        sender=sender,
        receiver=receiver
    )
    db.session.add(message)
    db.session.commit()

    # 构建消息对象
    message_data = {
        'id': message.id,
        'sender': sender,
        'receiver': receiver,
        'content': data['content'],
        'timestamp': datetime.datetime.now().isoformat()  # 使用本地时间而不是UTC
    }

    # 正确广播消息：发送给接收者和发送者自己
    emit('new_message', message_data, room=receiver)  # 发送给接收者
    # emit('new_message', message_data, room=sender)  # 发送给发送者（用于本地显示）


# def get_local_ip():
#     """获取本机局域网IP地址"""
#     try:
#         # 创建临时socket连接获取本机IP
#         s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
#         s.connect(("8.8.8.8", 80))  # 连接公共DNS
#         local_ip = s.getsockname()[0]
#         s.close()
#         return local_ip
#     except:
#         return '127.0.0.1'  # 失败时回退到localhost


if __name__ == '__main__':
    with app.app_context():
        db.create_all()

    # 自动获取IP并运行
    # host = get_local_ip()
    # port = 5001
    # socketio.run(app, host=host, port=port, debug=True)
    socketio.run(app, debug=True)
